#=

Migration and the Family testing Julia file

Reads in programs, estimates models, and performs counterfactuals while reporting/exporting results
Author: GA
=#

#add more cores!
using Distributed
addprocs(3)

@everywhere using Interpolations, Parameters, Distributions, DataFrames, CSV, LinearAlgebra, Statistics, Random, Profile, NLopt, SharedArrays
@everywhere include("icflm_setup.jl")
@everywhere include("icflm_valfunc.jl")
@everywhere include("icflm_estimation.jl")
@everywhere include("icflm_simulation.jl")
@everywhere include("icflm_readin.jl")

#####Change the subsequent line as appropriate to run the model
@everywhere dir = "C:\\Users\\garre\\Documents\\Work\\Papers\\Mig_Fam\\icflm_replication_archive\\Model"
@everywhere cd(dir)

#testing function that lets us avoid holding huge matrices in memory
function Test_simul(guess::Array{Float64,1}, dir::String; race::Int64 = 1)
    
    #load appropriate utilities
    package = Readin(dir)    
    if race == 2
        package = Readin_whites(dir)
    elseif race == 3
        package = Readin_blacks(dir)
    end
    estim_data, simul_data, trans_probs, div_chars = package[1], package[2], package[3], package[4]

    prim, est, res = Initialize(guess, trans_probs, div_chars) #initialize primitives, estimands, and results vectors
    Backward_Induct(prim, est, res; race = race) #compute all value functions
    likelihood, lhoods = Likelihood(prim, est, res, estim_data)
    simul_output = Simulate(prim, est, res, simul_data, lhoods; race = race)
    likelihood, lhoods, simul_output
end

######run for all races
for r = 1:1
    tempthing = CSV.read("$dir\\estimates_$r.csv", DataFrame; skipto=1)
    tempthing2 = Array{Float64,2}(tempthing)'
    tempthing3 = zeros(length(tempthing2))
    for iter = 1:length(tempthing2)
        tempthing3[iter] = tempthing2[iter]
    end
    guess_init = copy(tempthing3)
    
    ###run test stuff
    obj, lhoods, simul_output = Test_simul(guess_init, dir; race = r)
    println(obj)
    CSV.write("$dir\\simulated_data\\simulated_data_base_$r.csv", DataFrame(simul_output, :auto), header=false) #write CSV file
end

############COUNTERFACTUALS#############
#get basdeline type probabilities
function cfact_setup(guess::Array{Float64,1}, dir::String; race::Int64 = 1)
    
    #load appropriate utilities
    package = Readin(dir)    
    if race == 2
        package = Readin_whites(dir)
    elseif race == 3
        package = Readin_blacks(dir)
    end
    estim_data, simul_data, trans_probs, div_chars = package[1], package[2], package[3], package[4]
    
    prim, est, res = Initialize(guess, trans_probs, div_chars) #initialize primitives, estimands, and results vectors
    Backward_Induct(prim, est, res; race = race) #compute all value functions
    likelihood, lhoods = Likelihood(prim, est, res, estim_data)
    lhoods
end

#cfact simulation function
function Simulate_cfact(guess::Array{Float64,1}, dir::String, lhoods::Array{Float64, 2}; cfact::Int64=0, race::Int64 = 1)
    
    #load appropriate utilities
    package = Readin(dir)    
    if race == 2
        package = Readin_whites(dir)
    elseif race == 3
        package = Readin_blacks(dir)
    end
    estim_data, simul_data, trans_probs, div_chars = package[1], package[2], package[3], package[4]

    prim, est, res = Initialize(guess, trans_probs, div_chars) #initialize primitives, estimands, and results vectors
    Backward_Induct(prim, est, res; cfact = cfact, race = race) #compute all value functions
    simul_output = Simulate_lifecycle(prim, est, res, simul_data, lhoods; cfact = cfact, race = race)
    simul_output
end

#loop over races in counterfactuals
for r = 1:3
    tempthing = CSV.read("$dir\\estimates_$r.csv", DataFrame; skipto=1)
    tempthing2 = Array{Float64,2}(tempthing)'
    tempthing3 = zeros(length(tempthing2))
    for iter = 1:length(tempthing2)
        tempthing3[iter] = tempthing2[iter]
    end
    guess_init = copy(tempthing3)
    lhoods = cfact_setup(guess_init, dir; race = r)

    ######baseline#####
    simul_output_lifecycle = Simulate_cfact(guess_init, dir, lhoods; race = r)
    CSV.write("$dir\\simulated_data\\simulated_data_baseline_$r.csv", DataFrame(simul_output_lifecycle, :auto), header=false) #write CSV file

    #####grandparent cfacts#####
    ##get grandparent transfer everywhere
    simul_output_lifecycle = Simulate_cfact(guess_init, dir, lhoods; cfact = 1, race = r)
    CSV.write("$dir\\simulated_data\\simulated_data_cfact_nolock_$r.csv", DataFrame(simul_output_lifecycle, :auto), header=false) #write CSV file

    ##no grandparents
    guess_noparent = copy(guess_init)
    guess_noparent[11] = 0.0 #kill all grandparents!
    guess_noparent[12] = 0.0 #kill all grandparents!
    simul_output_lifecycle = Simulate_cfact(guess_noparent, dir, lhoods; race = r)
    CSV.write("$dir\\simulated_data\\simulated_data_cfact_noparent_$r.csv", DataFrame(simul_output_lifecycle, :auto), header=false) #write CSV file

    #####childcare policies#####
    ###national subsidy
    simul_output_lifecycle = Simulate_cfact(guess_init, dir, lhoods; cfact = 2, race = r)
    CSV.write("$dir\\simulated_data\\simulated_data_cfact_natsub_$r.csv", DataFrame(simul_output_lifecycle, :auto), header=false) #write CSV file

    ###local subsidy
    simul_output_lifecycle = Simulate_cfact(guess_init, dir, lhoods; cfact = 3, race = r)
    CSV.write("$dir\\simulated_data\\simulated_data_cfact_locsub_$r.csv", DataFrame(simul_output_lifecycle, :auto), header=false) #write CSV file

    ###halve childcare costs
    simul_output_lifecycle = Simulate_cfact(guess_init, dir, lhoods; cfact = 4, race = r)
    CSV.write("$dir\\simulated_data\\simulated_data_cfact_halfccc_$r.csv", DataFrame(simul_output_lifecycle, :auto), header=false) #write CSV file

    ###smaller reduction
    simul_output_lifecycle = Simulate_cfact(guess_init, dir, lhoods; cfact = 5, race = r)
    CSV.write("$dir\\simulated_data\\simulated_data_cfact_tweakccc_$r.csv", DataFrame(simul_output_lifecycle, :auto), header=false) #write CSV file    
end




##################################################